package com.xfusion.chatglm.data.domain.openai.service;

import cn.bugstack.chatglm.IOpenAiApi;
import cn.bugstack.chatglm.model.*;
import cn.bugstack.chatglm.session.OpenAiSession;
import com.alibaba.fastjson.JSON;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.xfusion.chatglm.data.domain.openai.model.aggregates.ChatProcessAggregate;
import com.xfusion.chatglm.data.domain.openai.model.entity.ChoiceEntity;
import com.xfusion.chatglm.data.domain.openai.model.entity.MessageEntity;
import com.xfusion.chatglm.data.domain.openai.model.entity.RuleLogicEntity;
import com.xfusion.chatglm.data.domain.openai.model.entity.UserAccountQuotaEntity;
import com.xfusion.chatglm.data.domain.openai.model.valobj.LogicCheckTypeVO;
import com.xfusion.chatglm.data.domain.openai.service.AbstractChatService;
import com.xfusion.chatglm.data.domain.openai.service.rule.ILogicFilter;
import com.xfusion.chatglm.data.domain.openai.service.rule.factory.DefaultLogicFactory;
import com.xfusion.chatglm.data.types.enums.ChatGLMModel;
import com.xfusion.chatglm.data.types.exception.ChatGLMException;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.stream.Collectors;

@Service
@Slf4j
public class ChatService extends AbstractChatService {

    @Resource
    protected OpenAiSession chatGlMOpenAiSession;
    @Resource
    private DefaultLogicFactory logicFactory;

    /**
     * 聊天响应
     *
     * 1.封装聊天过程中的用户消息作为请求参数
     * 2.向后端服务发起聊天完成请求，并注册事件监听器
     * 3.根据接收到的不同事件类型，实时推送聊天回复给客户端或记录聊天结束状态
     * 4.在事件源关闭时，通知客户端聊天响应流程结束
     *
     * @param responseBodyEmitter
     * @param chatProcess
     */
    @Override
    protected void doMessageResponse(ChatProcessAggregate chatProcess, ResponseBodyEmitter responseBodyEmitter) throws JsonProcessingException, InterruptedException {
        // 1. 请求消息
        //从chatprocess对象中获取用户发出的消息（对话历史）
        //这里用了流式处理，每个消息实体被映射到一个ChatCompletionRequest.Prompt对象中，其中role为user。content字段中存储了用户发出的消息内容。
        List<ChatCompletionRequest.Prompt> prompts = chatProcess.getMessages().stream()
                .map(entity -> ChatCompletionRequest.Prompt.builder()
                        .role(Role.user.getCode())
                        .content(entity.getContent())
                        .build())
                .collect(Collectors.toList());
        //最终被收集到一个列表prompts中，作为后续请求的输入

        // 2. 封装参数
        //构建一个chatcompletionrequest对象，设置模型、消息列表等参数，用于向后端服务发起聊天完成请求
        ChatCompletionRequest request = new ChatCompletionRequest();
        request.setModel(Model.valueOf(ChatGLMModel.get(chatProcess.getModel()).name())); // chatGLM_6b_SSE、chatglm_lite、chatglm_lite_32k、chatglm_std、chatglm_pro
        request.setPrompt(prompts);//对话历史

        //开启会话，注册一个EventSourceListener对象，用于接收后端服务返回的聊天结果
        chatGlMOpenAiSession.completions(request, new EventSourceListener() {
            /**
             * 当接收到一个事件时，onEvent方法会被调用，
             * @param eventSource   事件源
             * @param id    事件ID
             * @param type  事件类型
             * @param data  事件数据
             */
            @Override
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                log.info(data);
                //将数据反序列化为chatcompletionresponse对象,然后提取聊天回复内容response
                ChatCompletionResponse response = JSON.parseObject(data, ChatCompletionResponse.class);

                // 事件类型为add（增量更新）
                if (EventType.add.getCode().equals(type)){
                    try {
                        //将提取到的聊天回复内容发送给客户端，这意味着，每当后端生成新的消息，它都会被立即发送给客户端。
                        responseBodyEmitter.send(response.getData());
                    } catch (Exception e) {
                        throw new ChatGLMException(e.getMessage());
                    }
                }

                // type 消息类型，add 增量，finish 结束，error 错误，interrupted 中断
                if (EventType.finish.getCode().equals(type)) {
                    //同样反序列化，获取meta元数据，将元数据进一步反序列化为chatcompletionresponse.meta对象，并记录日志，表明聊天回复生成过程已经结束
                    ChatCompletionResponse.Meta meta = JSON.parseObject(response.getMeta(), ChatCompletionResponse.Meta.class);
                    log.info("[输出结束] Tokens {}", JSON.toJSONString(meta));
                }
            }

            @Override
            public void onClosed(EventSource eventSource) {
                responseBodyEmitter.complete();
            }

        });


    }

    /**
     * 这里我使用了工厂模式、策略模式、流式构建器设计思想，对聊天规则进行统一处理
     * @param chatProcessAggregate 聊天聚合体
     * @param userAccountQuotaEntity 用户账户
     * @param logics 过滤规则代号
     * @return
     * @throws Exception
     */

    @Override
    protected RuleLogicEntity<ChatProcessAggregate> doCheckLogic(ChatProcessAggregate chatProcessAggregate, UserAccountQuotaEntity userAccountQuotaEntity, String... logics) throws Exception {
        Map<String, ILogicFilter> logicFilterMap = logicFactory.openLogicFilter();
        RuleLogicEntity<ChatProcessAggregate>entity = null;
        for (String code:logics){
            //跳过无效的逻辑代码，如果当前的规律规则代码是null，则跳过，继续处理下一个逻辑代码
            if (DefaultLogicFactory.LogicModel.NULL.getCode().equals(code)) continue;
            //否则，从逻辑过滤器映射中根据当前逻辑代码获取对应的逻辑过滤器，并调用其过滤方法，传入聊天聚合体和用户账户信息，返回一个规则逻辑实体对象。
            entity = logicFilterMap.get(code).filter(chatProcessAggregate,userAccountQuotaEntity);
            //判断过滤结果是否成功，如果失败，则直接返回，否则继续处理下一个逻辑代码。
            if (!LogicCheckTypeVO.SUCCESS.equals(entity.getType()))return entity;
        }
        //如果所有的逻辑代码都通过检查，则返回一个成功的规则逻辑实体对象，否则返回一个失败的规则逻辑实体对象。
        return entity != null ? entity:RuleLogicEntity.<ChatProcessAggregate>builder()
                .type(LogicCheckTypeVO.SUCCESS).data(chatProcessAggregate).build();
    }

}
